#!/usr/bin/python2
# -*- coding: utf-8 -*-
# Name:     skel.py
# Author:   xiooli <xioooli[at]yahoo.com.cn>
# Version:  111129

import uuid, re, copy, sys

reg_conn = re.compile(r'^[0-9 \t\n]+$')
reg_end = re.compile(r'^[ \t]*[Mm][ \t]+[Ee][Nn][Dd][ \t\n]*$')

class Skel():
    def __init__(self, skel = []):
        self.skel = skel
        self.mol_header = '''[NO NAME]
  SKEL  1031110031302D
Created with skel.py'''

    def get_skel(self):
        return self.skel

    def set_skel(self, skel):
        self.skel = skel
        return self

    def skel2dic(self, skel):
        dic = {}
        for atom in skel:
            dic[atom['uuid']] = atom
        return dic

    def get_mol(self):
        if self.skel:
            return self.skel2mol(self.skel)

    def skel2mol(self, skel):
        skelobj = skel
        skeldic = self.skel2dic(skelobj)
        coords = []
        bonds = []
        for idx, atom in enumerate(skelobj):
            coords.append('%s %s  0  0  0  0  0  0  0  0  0  0  0  0' \
                    %(''.join([('%.4f' %i).rjust(10) for i in atom['coords']]), atom['element']))
            for b in atom['bonds']:
                if not skeldic.has_key(b[0]):
                    continue
                for bidx, a in enumerate(skelobj):
                    if a['uuid'] == b[0]:
                        break
                #print str(idx+1).rjust(3), str(bidx+1).rjust(3), str(b[1]).rjust(3)
                bonds.append('%s%s%s%s  0  0  0' %(str(idx+1).rjust(3), str(bidx+1).rjust(3), str(b[1]).rjust(3), str(b[2]).rjust(3)))
        desc = '%s%s  0  0  0  0            999 V2000' %(str(len(skelobj)).rjust(3), str(len(bonds)).rjust(3))
        return '%s\n%s\n%s\n%s\nM  END\n' %(self.mol_header, desc, '\n'.join(coords), '\n'.join(bonds))

    def link_skel(self, skel, pos):
        #print pos, self.skel2mol(skel)
        skelobj0 = copy.deepcopy(self.skel)
        skelobj1 = skel

        if not skelobj0 or not skelobj1:
            return []

        skeldic0 = self.skel2dic(skelobj0)
        skeldic1 = self.skel2dic(skelobj1)
        group_skels = []
        group_uuids = []
        link_node = []
        this_uuids = []
        this_link_node = []
        results = []

        for atom in skelobj1:
            if atom['element'] == '*':
                group_uuids.append(atom['uuid'])

        if group_uuids:
            for u in group_uuids:
                for atom in skelobj1:
                    if u in [i[0] for i in atom['bonds']]:
                        if not atom in link_node:
                            link_node.append(copy.deepcopy(atom))

        for atom in skelobj0:
            if atom['element'] == pos:
                this_uuids.append(atom['uuid'])

        if this_uuids:
            for u in this_uuids:
                for atom in skelobj0:
                    if u in [i[0] for i in atom['bonds']]:
                        if not atom in this_link_node:
                            this_link_node.append(copy.deepcopy(atom))

        # the linked node of group mark element *
        for idx1, node1 in enumerate(link_node):
            for idx0, node0 in enumerate(this_link_node):
                dic0 = copy.deepcopy(skeldic0)
                dic1 = copy.deepcopy(skeldic1)
                dic0[node0['uuid']]['bonds'].append((node1['uuid'], 1, 0))
                # delete the marked nodes
                dic0.pop(this_uuids[idx0])
                #dic1.pop(group_uuids[idx1])
                # replace the other marked nodes on group into hydrogen
                for id in group_uuids:
                    if dic1.has_key(id):
                        dic1.pop(id)
                        #dic1[idx]['element'] = 'H'

                results.append(dic0.values() + dic1.values())
        return results #[ self.skel2mol(r) for r in results ]

class Mol():
    def __init__(self, mol = ''):
        self.mol = mol
        self.atoms = []
        self.skelobj = Skel()
        if mol:
            self.atoms = self.mol2skel(self.mol)

    def get_mol(self):
        if self.atoms:
            self.mol = self.skelobj.set_skel(self.atoms).get_mol()
        return self.mol

    def set_mol(self, mol):
        self.mol = mol
        self.atoms = self.mol2skel(self.mol)
        return self

    def get_skel(self):
        if self.atoms:
            return self.atoms
        elif self.mol:
            return self.mol2skel(self.mol)
        return []

    def skel2dic(self, skel):
        dic = {}
        for atom in skel:
            dic[atom['uuid']] = atom
        return dic

    def mol2skel(self, mol):
        atoms = []
        try:
            for l in mol.replace('\r', '').replace('\t', ' ').split('\n')[4:]:
                lst = [i for i in l.split(' ') if i != '']
                if lst and ( not reg_conn.match(l) ) and ( not reg_end.match(l) ):
                    atom = {'element': lst[3],
                            'coords': [float(i) for i in lst[0:3]],
                            'uuid': str(uuid.uuid4()),
                            'bonds': []}
                    atoms.append(atom)
                elif lst and not reg_end.match(l):
                    atoms[int(lst[0])-1]['bonds'].append((atoms[int(lst[1])-1]['uuid'], lst[2], lst[3]))
            return atoms
        except Exception, e:
            #print e
            return []

    def uuid2idx(self, uuid):
        for idx, atom in enumerate(self.atoms):
            if atom['uuid'] == uuid:
                return idx

    def idx2uuid(self, idx):
        try:
            return self.atoms[idx]
        except:
            return ''

    def add_atom(self, element = 'C', coords = (0, 0, 0)):
        atom = {'element': element,
                'coords': coords,
                'uuid': str(uuid.uuid4()),
                'bonds': []}
        self.atoms.append(atom)
        return self.uuid2idx(atom['uuid'])

    def add_bond(self, atom_idx0, atom_idx1, bond_order = 1, stero_type = 0):
        try:
            self.atoms[atom_idx0]['bonds'].append(
                    (self.atoms[atom_idx1]['uuid'], str(bond_order), str(stero_type)) )
        except:
            pass
        return self

    def remove_atom(self, idx):
        try:
            uuid = self.atoms[idx]
            self.atoms.pop(idx)
            for a in self.atoms:
                new_bonds = []
                for b in a['bonds']:
                    if not uuid in b:
                        new_bonds.append(b)
                a['bonds'] = new_bonds
        except:
            pass

if __name__ == '__main__':
    from mol_utils import Mol2pic
    s0 = open('d:\\a.mol').read()
    s1 = open('d:\\ph.mol').read()
    m0 = Mol(s0)
    m1 = Mol(s1)
    m1.add_atom('I')
    m1.add_bond(0, -1)
    #m1.remove_atom(4)
    skel = Skel()
    skel.set_skel(m0.get_skel())
    #print Skel(m1.get_skel()).get_mol()
    sk = m1.get_skel()
    #print Mol2pic(skel.link_skel(sk, 'Ja')).to_pic(auto_layout = True)
    #sys.stderr.write(Mol2pic(Skel(skel.link_skel(sk, 'Ja')[1]).get_mol()).to_mol(True))
    m=Mol()
    m.add_atom('*')
    for i in range(5):
        m.add_atom('C')
        m.add_bond(i, -1)
    s=m.get_mol()
    print s
